#!/usr/bin/env python
# coding: utf-8

"""
baidu2labelme.py
create by cuiyu 2025-05-23
说明：
  * 将百度飞桨平台识别的井盖状态转换为 labelme 的 json 标注文件

用法:
  baidu2labelme.py

更新:
2025-05-30:
"""

import json 
from pathlib import Path
from loguru import logger
import datetime
import sys
from PIL import Image


_APP_FILE = Path(sys.argv[0])
_APP_ROOT = _APP_FILE.parent
_APP_NAME = _APP_FILE.stem

class Baidu2Labelme:
    def __init__(self, baidu_json_file: Path, img_root: Path):
        """
        将 baidu 飞桨平台识别的井盖状态转换为 labelme 的 json 标注文件
        输出的 labelme 的 json 文件位于 img_root 目录下，文件名为 id 值

        Args:
          :baidu_json_file: Path,  百度飞桨平台识别的 json 文件
          :img_root: Path,         图像文件根目录

        Returns: None
        """
        self.baidu_json_file = Path(baidu_json_file)
        self.img_root = Path(img_root)
        pass


    def _get_image_size(self, file_path):
        with Image.open(file_path) as img:  # 使用上下文管理器确保资源释放
            return img.size  # 返回 (width, height) 元组


    def run(self) -> bool:
        """
        将 baidu 飞桨平台识别的井盖状态转换为 labelme 的 json 标注文件
        """
        json_str = self.baidu_json_file.read_text(encoding="utf-8")
        try:
            json_dict: dict = json.loads(json_str)
        except:
            logger.error(f"解析 json 文件失败 {str(self.baidu_json_file)}")
            return False
        
        # 循环处理每个图像文件，并显示进度
        num = len(json_dict.items())
        n_len = len(str(num))
        for i, (k, v) in enumerate(json_dict.items()):
            n = i+1
            print(f'{n:{n_len}} / {num} {float(n)/num*100:.2f}%  imgid={k}', end='\r')
            self._write_labelme_json(v)
        print("")

        return True


    def _write_labelme_json(self, file_info: dict) -> bool:
        """
        将单个百度飞桨中的单个结果转换为 Labelme JSON 格式，并保存为JSON文件
        label <= alertTypeName
        points[] <= objects
        shape_type == rectangle

        参数:
            :file_info: dict 单个图像文件对应的标注信息

        返回:
            :bool: False=没有有效标注，True=有有效的标注
        """
        id_ = file_info["id"]
        json_file: Path = self.img_root / (str(id_)+".json")
        img_file = self.img_root / (str(id_)+".jpg")
        imgW, imgH = self._get_image_size(img_file)
        labelme_json = {
            "version": "4.5.6",
            "flags": {},
            "shapes": [],
            "imagePath": str(img_file.name),
            "imageData": None,
            "imageHeight": imgH,
            "imageWidth": imgW
        }

        cls_name = file_info["alertTypeName"]
        obj_num = file_info["objectCount"]
        objs = file_info["objects"]

        # 遍历每个 obj
        for i in range(obj_num):
            obj = objs[i]
            x1 = obj["x"]
            y1 = obj["y"]
            x2 = x1 + obj["w"]
            y2 = y1 + obj["h"]
            shape = {
                "label":  cls_name,  
                "points": [[x1, y1], [x2, y2]],        # 矩形形点列表
                "group_id": None,
                "shape_type": "rectangle",
                "flags": {}
            }
            labelme_json["shapes"].append(shape)
        
        if (len(labelme_json["shapes"]) == 0):
            return False
        
        json_str = json.dumps(labelme_json, indent=4, ensure_ascii=False)
        json_file.write_text(json_str, encoding="utf-8")
        return True


def _logInit() -> None:
    # logger.add(sys.stderr, format="{time} {level} {message}", level="INFO")
    time = datetime.datetime.today()
    # f'the time is {time:%Y-%m-%d (%a) %H:%M:%S}'   # datetime时间格式
    logger.add(f"{_APP_ROOT}/log/{_APP_NAME}-{time:%Y%m%d}.log", rotation="500 MB", encoding='utf-8', retention="10 days")  
    logger.add(f"{_APP_ROOT}/log/{_APP_NAME}-err-{time:%Y%m%d}.log", rotation="500 MB", encoding='utf-8', level='ERROR', retention="30 days")  

    # t = time.strftime("%Y%m%d") 
    # logger.add(f"{baseDir}/log-{t}.log", rotation="500 MB", encoding='utf-8', retention="10 days")  
    # logger.add(f"{baseDir}/err-{t}.log", rotation="500 MB", encoding='utf-8', level='ERROR', retention="30 days")  

def _main():
    baidu_json_file = Path("C:/aaby/work/01-dev/jgjc/_nogit/百度标注/百度标注-改了类别.json")
    img_root = Path("C:/aaby/netdisk/mydata/ai-data/dataset/priv/井盖病害/原始数据/井盖-百度原始数据-20250513/images")
    b2l = Baidu2Labelme(baidu_json_file, img_root)

    print("将百度飞桨平台识别的井盖状态转换为 labelme 的 json 标注文件：")
    b2l.run()
    print(".")


if __name__=="__main__":
    _logInit()
    _main()
